"""
根据双亲的基因型数据，筛选出双亲纯合差异位点
并根据标记的分离比进行卡方检验，设置p-value=0.05
"""
import sys
import math
import scipy.stats as stats

vcf_file_path= sys.argv[1]
f_code = sys.argv[2]
m_code = sys.argv[3]
output = sys.argv[4]
p_value = float(sys.argv[5])

file_handle = open(vcf_file_path, "r")

locus_pos=[]
locus_separate_p=[]

line_num = 0
columns = []
while True:
    line = file_handle.readline()
    if "#CHROM" in line:
        columns = line.strip().split("\t")
        break
    else:
        line_num+=1


f_index = columns.index(f_code)
m_index = columns.index(m_code)

# 打开文件
with open(vcf_file_path, 'r', encoding="utf-8") as file:
    for i, line in enumerate(file):
        if i < line_num:
            continue
        else:
            line_list = line.strip().split("\t")
            if line_list[f_index] in ["0|0", "1|1"] and line_list[m_index] in ["0|0", "1|1"] and line_list[f_index] != line_list[m_index]:
                # 计算标记的偏分离情况
                population_gts = []
                for index, v in enumerate(line_list):
                    if index > 9 and index not in [f_index, m_index]:
                        population_gts.append(v)

                gts_value = [1 if v == line_list[-1] else 0 for v in population_gts]
                # 计算卡方统计量和P值
                chi2, p, dof, ex = stats.chi2_contingency(
                    [[gts_value.count(0), gts_value.count(1)], [math.ceil(len(gts_value)/2), math.ceil(len(gts_value)/2)]]
                )
                if p >= p_value:
                    locus_pos.append([line_list[0], str(int(line_list[1]) - 1), line_list[1]])



with open(f"./{output}_locus.bed", "w") as file:
    for item in locus_pos:
        text = '\t'.join(item)
        file.write(f"{text}\n")
